#include "kernel_operator.h"
#include <cfloat>
using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2;

class KernelPowsBF {
    public:
        __aicore__ inline KernelPowsBF() {}
        __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, uint32_t size,
                                    uint32_t smallCoreDataNum,
                                    uint32_t finalSmallTileNum, 
                                    uint32_t tileDataNum,
                                    uint32_t smallTailDataNum,
                                    uint32_t tailBlockNum) {
            ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!");
            this->tileDataNum = tileDataNum;

            this->coreDataNum = smallCoreDataNum; 
            this->tileNum = finalSmallTileNum;    
            this->tailDataNum = smallTailDataNum;  
    
            x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x1, this->coreDataNum);
            x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2 *)x2, this->coreDataNum);
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y, this->coreDataNum);
    
            pipe.InitBuffer(inQueueX1, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X1));
            pipe.InitBuffer(inQueueX2, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X2));
            pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_Y));

            pipe.InitBuffer(tmpBuffer3, this->tileDataNum * sizeof(float));
            pipe.InitBuffer(tmpBuffer4, this->tileDataNum * sizeof(float));

            pipe.InitBuffer(tmpBuffer5, this->tileDataNum * sizeof(float));
            pipe.InitBuffer(tmpBuffer6, this->tileDataNum * sizeof(float));
        }
    
        __aicore__ inline void process() {
            int32_t loopCount = this->tileNum;
            this->processDataNum = this->tileDataNum;
            for (int32_t i = 0; i < loopCount; i++)
            {
                if (i == this->tileNum - 1) {
                    this->processDataNum = this->tailDataNum;
                }
                CopyIn(i);
                Compute(i);
                CopyOut(i);
            }
        }
    private:
    
        __aicore__ inline void CopyIn(int32_t progress)
        {
            LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
            LocalTensor<DTYPE_X2> x2Local = inQueueX2.AllocTensor<DTYPE_X2>();
            
            DataCopy(x1Local, x1Gm[progress * this->tileDataNum], this->processDataNum);
            DataCopy(x2Local, x2Gm[progress * this->tileDataNum], this->processDataNum);
            
            inQueueX1.EnQue(x1Local);
            inQueueX2.EnQue(x2Local);
        }
    
        __aicore__ inline void Compute(int32_t progress)
        {
            LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
            LocalTensor<DTYPE_X2> x2Local = inQueueX2.DeQue<DTYPE_X2>();
            LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();

            LocalTensor<float> tmpTensor3 = tmpBuffer3.Get<float>();
            LocalTensor<float> tmpTensor4 = tmpBuffer4.Get<float>();
            LocalTensor<float> tmpTensor5 = tmpBuffer3.Get<float>();
            LocalTensor<float> tmpTensor6 = tmpBuffer4.Get<float>();

            Cast(tmpTensor3, x1Local, RoundMode::CAST_NONE, this->processDataNum);      //x1: bf16 -> float32
            Cast(tmpTensor4, x2Local, RoundMode::CAST_NONE, this->processDataNum);      //x2: bf16 -> float32

            Ln(tmpTensor5, tmpTensor3, this->processDataNum);      //float: ln(x1)
            Mul(tmpTensor6, tmpTensor5, tmpTensor4, this->processDataNum);     //x2 * Ln(x1)
            Exp(tmpTensor5, tmpTensor6, this->processDataNum);      //e^(x2 * ln(x1))

            Cast(yLocal, tmpTensor5, RoundMode::CAST_ROUND, this->processDataNum);        //y: float32 -> bf16
    
            outQueueY.EnQue(yLocal);
            inQueueX1.FreeTensor(x1Local);
            inQueueX2.FreeTensor(x2Local);
            
        }
    
        __aicore__ inline void CopyOut(int32_t progress) {
            LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
            
            DataCopy(yGm[progress * this->tileDataNum], yLocal, this->processDataNum);
            
            outQueueY.FreeTensor(yLocal);
        }
    
        TPipe pipe;
        TQue<TPosition::VECIN, BUFFER_NUM> inQueueX1, inQueueX2;
        TQue<TPosition::VECOUT, BUFFER_NUM> outQueueY;
        TBuf<TPosition::VECCALC> tmpBuffer3,tmpBuffer4,tmpBuffer5,tmpBuffer6;
    
        GlobalTensor<DTYPE_X1> x1Gm;
        GlobalTensor<DTYPE_X2> x2Gm;
        GlobalTensor<DTYPE_Y> yGm;
    
        uint32_t smallCoreDataNum;
        uint32_t tileDataNum;
        uint32_t coreDataNum;
        uint32_t tileNum;
        uint32_t tailDataNum;
        uint32_t processDataNum;
        
    
    };



class KernelPowsF16 {
    public:
        __aicore__ inline KernelPowsF16() {}
        __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, uint32_t size,
                                    uint32_t smallCoreDataNum,
                                    uint32_t finalSmallTileNum, 
                                    uint32_t tileDataNum,
                                    uint32_t smallTailDataNum,
                                    uint32_t tailBlockNum) {
            ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!");
            this->tileDataNum = tileDataNum;

            this->coreDataNum = smallCoreDataNum; 
            this->tileNum = finalSmallTileNum;    
            this->tailDataNum = smallTailDataNum;  
    
            x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x1, this->coreDataNum);
            x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2 *)x2, this->coreDataNum);
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y, this->coreDataNum);
    
            pipe.InitBuffer(inQueueX1, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X1));
            pipe.InitBuffer(inQueueX2, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X2));
            pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_Y));

            pipe.InitBuffer(tmpBuffer3, this->tileDataNum * sizeof(float));
            pipe.InitBuffer(tmpBuffer4, this->tileDataNum * sizeof(float));

            pipe.InitBuffer(tmpBuffer5, this->tileDataNum * sizeof(float));
            pipe.InitBuffer(tmpBuffer6, this->tileDataNum * sizeof(float));
        }
    
        __aicore__ inline void process() {
            int32_t loopCount = this->tileNum;
            this->processDataNum = this->tileDataNum;
            for (int32_t i = 0; i < loopCount; i++)
            {
                if (i == this->tileNum - 1) {
                    this->processDataNum = this->tailDataNum;
                }
                CopyIn(i);
                Compute(i);
                CopyOut(i);
            }
        }
    private:
    
        __aicore__ inline void CopyIn(int32_t progress)
        {
            LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
            LocalTensor<DTYPE_X2> x2Local = inQueueX2.AllocTensor<DTYPE_X2>();
            
            DataCopy(x1Local, x1Gm[progress * this->tileDataNum], this->processDataNum);
            DataCopy(x2Local, x2Gm[progress * this->tileDataNum], this->processDataNum);
            
            inQueueX1.EnQue(x1Local);
            inQueueX2.EnQue(x2Local);
        }
    
        __aicore__ inline void Compute(int32_t progress)
        {
            LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
            LocalTensor<DTYPE_X2> x2Local = inQueueX2.DeQue<DTYPE_X2>();
            LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();

            LocalTensor<float> tmpTensor3 = tmpBuffer3.Get<float>();
            LocalTensor<float> tmpTensor4 = tmpBuffer4.Get<float>();
            LocalTensor<float> tmpTensor5 = tmpBuffer3.Get<float>();
            LocalTensor<float> tmpTensor6 = tmpBuffer4.Get<float>();

            Cast(tmpTensor3, x1Local, RoundMode::CAST_NONE, this->processDataNum);      //x1: bf16 -> float32
            Cast(tmpTensor4, x2Local, RoundMode::CAST_NONE, this->processDataNum);      //x2: bf16 -> float32

            Ln(tmpTensor5, tmpTensor3, this->processDataNum);      //float: ln(x1)
            Mul(tmpTensor6, tmpTensor5, tmpTensor4, this->processDataNum);     //x2 * Ln(x1)
            Exp(tmpTensor5, tmpTensor6, this->processDataNum);      //e^(x2 * ln(x1))

            Cast(yLocal, tmpTensor5, RoundMode::CAST_NONE, this->processDataNum);        //y: float32 -> bf16

            outQueueY.EnQue(yLocal);
            inQueueX1.FreeTensor(x1Local);
            inQueueX2.FreeTensor(x2Local);

        }
    
        __aicore__ inline void CopyOut(int32_t progress) {
            LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
            
            DataCopy(yGm[progress * this->tileDataNum], yLocal, this->processDataNum);
            
            outQueueY.FreeTensor(yLocal);
        }
    
        TPipe pipe;
        TQue<TPosition::VECIN, BUFFER_NUM> inQueueX1, inQueueX2;
        TQue<TPosition::VECOUT, BUFFER_NUM> outQueueY;
        TBuf<TPosition::VECCALC> tmpBuffer3,tmpBuffer4,tmpBuffer5,tmpBuffer6;
    
        GlobalTensor<DTYPE_X1> x1Gm;
        GlobalTensor<DTYPE_X2> x2Gm;
        GlobalTensor<DTYPE_Y> yGm;
    
        uint32_t smallCoreDataNum;
        uint32_t tileDataNum;
        uint32_t coreDataNum;
        uint32_t tileNum;
        uint32_t tailDataNum;
        uint32_t processDataNum;
    };



class KernelPowsFloat {
        public:
            __aicore__ inline KernelPowsFloat() {}
            __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, uint32_t size,
                                        uint32_t smallCoreDataNum,
                                        uint32_t finalSmallTileNum, 
                                        uint32_t tileDataNum,
                                        uint32_t smallTailDataNum,
                                        uint32_t tailBlockNum) {
                ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!");
                this->tileDataNum = tileDataNum;
        
                this->coreDataNum = smallCoreDataNum; 
                this->tileNum = finalSmallTileNum;    
                this->tailDataNum = smallTailDataNum;  
        
                x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x1, this->coreDataNum);
                x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2 *)x2, this->coreDataNum);
                yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y, this->coreDataNum);
        
                pipe.InitBuffer(inQueueX1, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X1));
                pipe.InitBuffer(inQueueX2, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X2));
                pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_Y));
        
                pipe.InitBuffer(tmpBuffer1, this->tileDataNum * sizeof(DTYPE_X1));
                pipe.InitBuffer(tmpBuffer2, this->tileDataNum * sizeof(DTYPE_X1));
            }
        
            __aicore__ inline void process() {
                int32_t loopCount = this->tileNum;
                this->processDataNum = this->tileDataNum;
                for (int32_t i = 0; i < loopCount; i++)
                {
                    if (i == this->tileNum - 1) {
                        this->processDataNum = this->tailDataNum;
                    }
                    CopyIn(i);
                    Compute(i);
                    CopyOut(i);
                }
            }
        private:
        
            __aicore__ inline void CopyIn(int32_t progress)
            {
                LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
                LocalTensor<DTYPE_X2> x2Local = inQueueX2.AllocTensor<DTYPE_X2>();
                
                DataCopy(x1Local, x1Gm[progress * this->tileDataNum], this->processDataNum);
                DataCopy(x2Local, x2Gm[progress * this->tileDataNum], this->processDataNum);
                
                inQueueX1.EnQue(x1Local);
                inQueueX2.EnQue(x2Local);
            }
        
            __aicore__ inline void Compute(int32_t progress)
            {
                LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
                LocalTensor<DTYPE_X2> x2Local = inQueueX2.DeQue<DTYPE_X2>();
                LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();

                LocalTensor<DTYPE_X1> tmpTensor1 = tmpBuffer1.Get<DTYPE_X1>();
                LocalTensor<DTYPE_X1> tmpTensor2 = tmpBuffer2.Get<DTYPE_X1>();
                
                Ln(tmpTensor1, x1Local, this->processDataNum);      //ln(x1)
                Mul(tmpTensor2, tmpTensor1, x2Local, this->processDataNum);     //x2 * Ln(x1)
                Exp(yLocal, tmpTensor2, this->processDataNum);      //e^(x2 * ln(x1))
        
                outQueueY.EnQue(yLocal);
                inQueueX1.FreeTensor(x1Local);
                inQueueX2.FreeTensor(x2Local);
            }
        
            __aicore__ inline void CopyOut(int32_t progress) {
                LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
                
                DataCopy(yGm[progress * this->tileDataNum], yLocal, this->processDataNum);
                
                outQueueY.FreeTensor(yLocal);
            }
        
            TPipe pipe;
            TQue<TPosition::VECIN, BUFFER_NUM> inQueueX1, inQueueX2;
            TQue<TPosition::VECOUT, BUFFER_NUM> outQueueY;
            TBuf<TPosition::VECCALC> tmpBuffer1,tmpBuffer2;
        
            GlobalTensor<DTYPE_X1> x1Gm;
            GlobalTensor<DTYPE_X2> x2Gm;
            GlobalTensor<DTYPE_Y> yGm;
        
            uint32_t smallCoreDataNum;
            uint32_t tileDataNum;
            uint32_t coreDataNum;
            uint32_t tileNum;
            uint32_t tailDataNum;
            uint32_t processDataNum;
        };
    


class KernelPowsBC {
    public:
        __aicore__ inline KernelPowsBC() {}
        __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, uint32_t size,
                                    uint32_t smallCoreDataNum,
                                    uint32_t finalSmallTileNum, 
                                    uint32_t tileDataNum,
                                    uint32_t smallTailDataNum,
                                    uint32_t tailBlockNum, BroadCastTiling BCT) {
            ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!");
            this->tileDataNum = tileDataNum;

            this->coreDataNum = smallCoreDataNum; 
            this->tileNum = finalSmallTileNum;    
            this->tailDataNum = smallTailDataNum;  

            this->dimNum = BCT.y_dimNum;
            this->y_dims = BCT.y_dims;
            this->x1_dims = BCT.x1_dims;
            this->x2_dims = BCT.x2_dims;
            this->y_sumDims = BCT.y_sumDims;
            this->x1_sumDims = BCT.x1_sumDims;
            this->x2_sumDims = BCT.x2_sumDims;
    
            x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x1, this->coreDataNum);   
            x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2 *)x2, this->coreDataNum);
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y, this->coreDataNum);

            pipe.InitBuffer(tmpBuffer1, sizeof(float));
            pipe.InitBuffer(tmpBuffer2, sizeof(float));
            pipe.InitBuffer(tmpBuffer3, sizeof(float));
            pipe.InitBuffer(tmpBuffer4, sizeof(float));
            pipe.InitBuffer(tmpBuffer5, sizeof(float));
            pipe.InitBuffer(tmpBuffer6, sizeof(bfloat16_t));
        }
    
        __aicore__ inline void process() {
            
            for(uint32_t i=0; i<this->coreDataNum; i++){
                uint32_t x1_pos = 0, x2_pos = 0;
                for(uint32_t j=0; j<this->dimNum; j++)
                {
                    if(this->x1_dims[j] != 1)
                    {
                        x1_pos += this->x1_sumDims[j] * (i / this->y_sumDims[j] % this->y_dims[j]);
                    }
                    if(this->x2_dims[j] != 1)
                    {
                        x2_pos += this->x2_sumDims[j] * (i / this->y_sumDims[j] % this->y_dims[j]);
                    }
                }

                LocalTensor<float> tmpTensor1 = tmpBuffer1.Get<float>();
                LocalTensor<float> tmpTensor2 = tmpBuffer2.Get<float>();
                LocalTensor<float> tmpTensor3 = tmpBuffer3.Get<float>();
                LocalTensor<float> tmpTensor4 = tmpBuffer4.Get<float>();
                LocalTensor<float> tmpTensor5 = tmpBuffer5.Get<float>();
                LocalTensor<bfloat16_t> tmpTensor6 = tmpBuffer6.Get<bfloat16_t>();

                auto x1_val = x1Gm.GetValue(x1_pos);  // 获取值
                auto x2_val = x2Gm.GetValue(x2_pos);
                bfloat16_t bf16_dummy;
                float16_t f16_dummy;
                float32_t x1,x2;

                // if constexpr (std::is_same_v<DTYPE_X1, bfloat16_t>)
                if (GetDataType(x1_val) == GetDataType<bfloat16_t>(bf16_dummy)) {  // 检查是否是 half（可能是 bfloat16 或 float16）
                    float32_t x1 = max(ToFloat(x1_val), FLT_MIN);
                    float32_t x2 = max(ToFloat(x2_val), FLT_MIN);
                    
                    Duplicate(tmpTensor1, x1, (uint8_t)1);   //bf16的数据变成tensor
                    Duplicate(tmpTensor2, x2, (uint8_t)1);

                    Ln(tmpTensor3, tmpTensor1, (uint8_t)1);      //ln(x1)
                    Mul(tmpTensor4, tmpTensor3, tmpTensor2, (uint8_t)1);     //x2 * Ln(x1)
                    Exp(tmpTensor5, tmpTensor4, (uint8_t)1);      //e^(x2 * ln(x1))

                    Cast(tmpTensor6, tmpTensor5, RoundMode::CAST_RINT, 1);

                    yGm.SetValue(i, tmpTensor6.GetValue(0));
                }else {
                    x1 = x1_val;
                    x2 = x2_val;
                    Duplicate(tmpTensor1, x1, (uint8_t)1);
                    Duplicate(tmpTensor2, x2, (uint8_t)1);

                    Ln(tmpTensor3, tmpTensor1, (uint8_t)1);      //ln(x1)
                    Mul(tmpTensor4, tmpTensor3, tmpTensor2, (uint8_t)1);     //x2 * Ln(x1)
                    Exp(tmpTensor5, tmpTensor4, (uint8_t)1);      //e^(x2 * ln(x1))

                    yGm.SetValue(i, tmpTensor5.GetValue(0));
                }
            }
        }
        private:    
        TPipe pipe;
        
        GlobalTensor<DTYPE_X1> x1Gm;
        GlobalTensor<DTYPE_X2> x2Gm;
        GlobalTensor<DTYPE_Y> yGm;
        TBuf<TPosition::VECCALC> tmpBuffer1,tmpBuffer2,tmpBuffer3,tmpBuffer4,tmpBuffer5,tmpBuffer6,tmpBuffer7,tmpBuffer8;
    
        uint32_t smallCoreDataNum;
        uint32_t tileDataNum;
        uint32_t coreDataNum;
        uint32_t tileNum;
        uint32_t tailDataNum;
        uint32_t processDataNum;
        
        int32_t dimNum;
        int32_t *y_dims, *y_sumDims;
        int32_t *x1_dims, *x1_sumDims;
        int32_t *x2_dims, *x2_sumDims;
    
};


    
extern "C" __global__ __aicore__ void pows(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    if (tiling_data.tilingKey == 1) {
        KernelPowsBF op;
        op.Init(x1, x2, y, tiling_data.size, 
            tiling_data.smallCoreDataNum,
            tiling_data.finalSmallTileNum,
            tiling_data.tileDataNum,
            tiling_data.smallTailDataNum,
            tiling_data.tailBlockNum);
        op.process();
    }else if (tiling_data.tilingKey == 2) {
        KernelPowsF16 op;
        op.Init(x1, x2, y, tiling_data.size, 
            tiling_data.smallCoreDataNum,
            tiling_data.finalSmallTileNum,
            tiling_data.tileDataNum,
            tiling_data.smallTailDataNum,
            tiling_data.tailBlockNum);
        op.process();
    } else if (tiling_data.tilingKey == 3) {
        KernelPowsFloat op;
        op.Init(x1, x2, y, tiling_data.size, 
            tiling_data.smallCoreDataNum,
            tiling_data.finalSmallTileNum,
            tiling_data.tileDataNum,
            tiling_data.smallTailDataNum,
            tiling_data.tailBlockNum);
        op.process();
    }else if (tiling_data.tilingKey == 4) {
        KernelPowsBC op;
        op.Init(x1, x2, y, tiling_data.size, 
            tiling_data.smallCoreDataNum,
            tiling_data.finalSmallTileNum,
            tiling_data.tileDataNum,
            tiling_data.smallTailDataNum,
            tiling_data.tailBlockNum,
            tiling_data.BCT);
        op.process();
    }else {
        printf("Error: Tiling key is not valid!\n");
    }
}